"""
PySCeS - Core2
Author: Brett G. Olivier (brettolivier@gmail.com)
Copyright (C) Brett. G. Olivier, Stellenbosch, 2007.
The source code contained in this file is not licensed for use
or redistribution without the author's prior and express permission.
"""

from version import current_version
__version__ = current_version()

__doc__ = '''PySCeS NewCore module: useful for exploring model component relations'''

import os
CurrentDirectory = os.path.dirname(os.path.abspath(os.sys.argv[0]))

import numpy
##  import scipy

from InfixParser import MyInfixParser

InfixParser = MyInfixParser()
InfixParser.buildlexer()
InfixParser.buildparser(debug=0, debugfile='infix.dbg', tabmodule='infix_tabmodule')
InfixParser.setNameStr('self.', '()')

class MapList(list):
    def __init__(self, *args):
        list.__init__(self,*args)
        
    def asSet(self):
        return set(self.__getslice__(0, self.__len__()))

class NewCoreBase(object):
    name = None

    def getName(self):
        return self.name

    def setName(self,name):
        self.name = name

    def get(self, attr):
        """Return an attribute whose name is str(attr)"""
        try:
            return self.__getattribute__(attr)
        except:
            print "%s is not an attribute of this instance" % attr
            return None

class NumberBase(NewCoreBase):
    value = None
    value_initial = None
    
    def __call__(self):
        return self.value
    
    def getValue(self):
        return self.value
        
    def setValue(self, v):
        self.value = v

class Species(NumberBase):
    subs = None
    prods = None
    mods = None
    fixed = False
    compartment = None

    def __init__(self, name, value):
        self.setName(name)
        self.value = value
        self.value_initial = value
        self.subs = []
        self.prods = []
        self.mods = []

    def setCompartment(self, c):
        self.compartment = c

    def setSubstrate(self, reaction):
        self.__setattr__(reaction.name, reaction)
        self.subs.append(reaction)

    def setProduct(self, reaction):
        self.__setattr__(reaction.name, reaction)
        self.prods.append(reaction)

    def setModifier(self, reaction):
        self.__setattr__(reaction.name, reaction)
        self.mods.append(reaction)

    def isSubstrateOf(self):
        return MapList([r.name for r in self.subs])

    def isProductOf(self):
        return MapList([r.name for r in self.prods])

    def isModifierOf(self):
        return MapList([r.name for r in self.mods])

    def isReagentOf(self):
        return MapList(self.isSubstrateOf() + self.isProductOf())

class Reaction(NewCoreBase):
    modifiers = None
    substrates = None
    products = None
    stoichiometry = None
    parameters = None
    reversible = True
    formula = None
    code_string = None
    rate = None
    xcode = None
    _names = None

    def __call__(self):
        ##  exec(self.code_string)
        exec(self.xcode)
        return self.rate

    def __init__(self, name):
        self.setName(name)
        self.modifiers = []
        self.substrates = []
        self.products = []
        self.stoichiometry = {}
        self.parameters = []

    def addSubstrate(self, species):
        self.__setattr__(species.name, species)
        self.substrates.append(species)

    def addProduct(self, species):
        self.__setattr__(species.name, species)
        self.products.append(species)

    def addModifier(self, species):
        self.__setattr__(species.name, species)
        self.modifiers.append(species)

    def addFormula(self, formula):
        self.formula = formula
        InfixParser.parse(formula)
        self._names = InfixParser.names
        #print self._names
        self.code_string = 'self.rate=%s' % InfixParser.output
        self.xcode = compile(self.code_string, '<string>', 'exec')

    def addParameter(self, par):
        self.__setattr__(par.name, par)
        self.parameters.append(par)        

    def hasProducts(self, t=type):
        return MapList([p.name for p in self.products])

    def hasSubstrates(self):
        return MapList([s.name for s in self.substrates])

    def hasModifiers(self):
        return MapList([m.name for m in self.modifiers])

    def hasParameters(self):
        return MapList([p.name for p in self.parameters])

    def hasReagents(self):
        return MapList(self.hasSubstrates() + self.hasProducts())


class Parameter(NumberBase):
    association = None
    
    def __init__(self, name, value):
        self.name = name
        self.value = value
        self.value_initial = value
        self.association = []

    def setAssociation(self, reac):
        self.association.append(reac)
        self.__setattr__(reac.name, reac)

    def isParameterOf(self):
        return MapList([a.name for a in self.association])

class AssignmentRule(Parameter):
    formula = None
    code_string = None
    _names = None
    
    def __init__(self, name, value):
        Parameter.__init__(self, name, value)

    def __call__(self):
        exec(self.xcode)
        ##  exec(self.code_string)
        return self.value

    def addFormula(self, formula):
        self.formula = formula
        InfixParser.parse(formula)
        self.code_string = 'self.value=%s' % InfixParser.output
        self._names = InfixParser.names
        self.xcode = compile(self.code_string, '<string>', 'exec')
        
    def addModelAttr(self, obj):
        self.__setattr__(obj.name, obj)

class ODE(NewCoreBase):
    sdot = None
    value = None
    coefficients = None
    reactions = None
    independent = None
    ode_terms = None
    formula = ''
    formula_alt = ''
    code_string = 'self.value='
    code_string_alt = 'sdot='
    
    
    def __init__(self, species, independent=True):
        self.sdot = species
        self.name = 'ODE_'+species.name
        self.reactions = []
        self.coefficients = []
        self.ode_terms = []
        self.independent = independent
                
    def __call__(self):
        exec(self.code_string)
        return self.value

    def addReaction(self, reaction, coefficient):
        self.reactions.append(reaction)
        self.coefficients.append(coefficient)
        if coefficient > 0.0:
            if coefficient == 1.0:
                term = '+self.%s() ' % (reaction.name)
                aterm = '+(%s) ' % (reaction.code_string.replace('self.rate=','')) 
                fterm = '+%s' % (reaction.name)
                afterm = '+ (%s) ' % (reaction.formula) 
            else:
                term = '+%f*self.%s() ' % (abs(coefficient), reaction.name)
                aterm = '+%f*(%s) ' % (abs(coefficient), reaction.code_string.replace('self.rate=','')) 
                fterm = '+%f*%s' % (abs(coefficient), reaction.name)
                afterm = '+ %f*(%s) ' % (abs(coefficient), reaction.formula) 
        else:
            if coefficient == -1.0:
                term = '-self.%s() ' % (reaction.name)
                aterm = '-(%s) ' % (reaction.code_string.replace('self.rate=','')) 
                fterm = '-%s' % (reaction.name)
                afterm = '- (%s) ' % (reaction.formula) 
            else:
                term = '-%f*self.%s() ' % (abs(coefficient), reaction.name)
                aterm = '-%f*(%s) ' % (abs(coefficient), reaction.code_string.replace('self.rate=','')) 
                fterm = '-%f*%s' % (abs(coefficient), reaction.name)
                afterm = '- %f*(%s) ' % (abs(coefficient), reaction.formula) 
        self.ode_terms.append(term)
        self.code_string += term
        self.code_string_alt += aterm
        self.formula += fterm
        self.formula_alt += afterm
        self.__setattr__(reaction.name, reaction)
        ##  for obj in reaction.parameters+reaction.substrates+reaction.products+reaction.modifiers:
            ##  self.__setattr__(obj.name, obj)
    
    def hasReactions(self):
        return MapList([r.name for r in self.reactions])
        
    def getFormula(self):
        return self.code_string
        
    def getGlobalFormula(self):
        return self.code_string_alt
        
    
class StructMatrix(NewCoreBase):
    """
    This class is specifically designed to store structural matrix information
    give it an array and row/col index permutations it can generate its own
    row/col labels given the label src. It might become a more generic class
    but then row/col labels must be externally assigned -- 50/50 atm.
    """
    
    array = None
    ridx = None
    cidx = None
    row = None
    col = None

    def __init__(self, array, ridx, cidx, row=None, col=None):
        """
        Instantiate with array and matching row/col index arrays, optional label arrays
        """
        self.array = array
        self.ridx = ridx
        self.cidx = cidx
        self.row = row
        self.col = col
    
    def __call__(self):
        return self.array
        
    def getRowsByIdx(self, *args):
        """Return the rows referenced by index (1,3,5)"""
        return self.array.take(args, axis=0)
        
    def getColsByIdx(self, *args):
        """Return the columns referenced by index (1,3,5)"""
        return self.array.take(args, axis=1)

    def setRow(self, src):
        """
        Assuming that the row index array is a permutation (full/subset)
        of a source label array by supplying that source to setRow it 
        maps the row labels to ridx and creates self.row (row label list)
        """
        self.row = [src[r] for r in self.ridx]

    def setCol(self, src):
        """
        Assuming that the col index array is a permutation (full/subset)
        of a source label array by supplying that src to setCol 
        maps the row labels to cidx and creates self.col (col label list)
        """
        
        self.col = [src[c] for c in self.cidx]

    def getRowsByName(self, *args):
        """Return the rows referenced by label ('s','x','d')"""
        assert self.row != None, "\nI need row labels"
        try:
            return self.array.take([self.row.index(l) for l in args], axis=0)
        except Exception, ex:
            print ex
            print "\nValid row labels are: %s" % self.row
            return None
        
    def getColsByName(self, *args):
        """Return the columns referenced by label ('s','x','d')"""
        assert self.col != None, "\nI need column labels"
        try:
            return self.array.take([self.col.index(l) for l in args], axis=1)
        except Exception, ex:
            print ex
            print "Valid column labels are: %s" % self.col
            return None

    def getLabels(self, axis='all'):
        """Return the matrix labels ([rows],[cols]) where axis='row'/'col'/'all'"""
        if axis == 'row': return self.row
        elif axis == 'col': return self.col
        else: return self.row, self.col

    def getIndexes(self, axis='all'):
        """Return the matrix indexes ([rows],[cols]) where axis='row'/'col'/'all'"""
        if axis == 'row': return self.ridx
        elif axis == 'col': return self.cidx
        else: return self.ridx, self.cidx

class NewCore(NewCoreBase):
    ndic = None
    reactions = None
    species = None
    species_variable = None
    __model__ = None
    __InitStrings__ = None
    __InitDict__ = None
    __not_inited__ = None
    global_parameters = None
    __parameter_store__ = None
    forcing_functions = None
    __forcing_function_dic__ = None
    __AssignmentRuleClass__ = None
    # new
    compartments = None
    description = "Pysces Core 2: pysXgen"
    global_units = None
    stoichiometric_matrix = None
    struct = None
    ODEs = None

    def __init__(self, model):
        self.setName(model.ModelFile[:-4])
        self.ndic = model._PysMod__NetworkDict
        self.__model__ = model
        self.__InitStrings__ = [s.replace('self.','') for s in self.__model__._PysMod__InitStrings]
        self.__InitDict__ = {}
        # get initialisation strings
        for ii in self.__InitStrings__:
            l,r = ii.split('=')
            self.__InitDict__.setdefault(self.__cleanString__(l), float(self.__cleanString__(r)))
        # set to current model value
        for k in self.__InitDict__.keys():
            self.__InitDict__[k] = getattr(self.__model__, k) 
        self.global_parameters = []
        self.__parameter_store__ = []
        self.__not_inited__ = []
        self.__function_forced_str__ = model._Function_forced
        self.forcing_functions = []
        self.__forcing_function_dic__ = {}
        self.compartments = [('Cell', 1.0)]
        self.global_units = {'mole':(1, 0),
                             'litre':(1, 0),
                             'second':(1, 0)}
        
        try:
            self.parseForcingFunctions()
        except Exception, ex:
            print 'No forcing functions (%s).' % ex
        self.addSpecies()
        self.addReactions()
        self.generateMappings()
        self.setAssignmentRules()
        assert len(self.__not_inited__) < 1, "\nERROR: Uninitialised parameters: %s" % self.__not_inited__
        
    def __cleanString__(self,s):
        s = s.lstrip()
        s = s.rstrip()
        return s

    def parseForcingFunctions(self):
        self.__forcing_function_dic__ = {}
        ff = self.__function_forced_str__.split('\n')
        for f in ff:
            if f != '':
                f =  f.split('=')
                f[0] = f[0].replace('self.','')
                f[1] = f[1].replace('self.','')
                self.__forcing_function_dic__.setdefault(self.__cleanString__(f[0]), self.__cleanString__(f[1]))

    def setDescription(self, txt):
        self.description = str(txt)
        
    def getDescription(self):
        return str(self.description)

    def setGlobalUnits(self, **kwargs):
        for un in kwargs.keys():
            self.global_units[un] = (kwargs[un][0], kwargs[un][1])
            print "Modified \"%s\" to be %i*%s*10**%i" % (un, kwargs[un][0], un, kwargs[un][1])

    def getGlobalUnits(self):
        return self.global_units

    def addOneSpecies(self, species, value, fix=False, comp=None):
        if comp == None: comp = self.compartments[0][0]
        s = Species(species, value)
        s.setCompartment(comp)        
        if fix: s.fixed = True
        self.__setattr__(species, s)
        self.species.append(s)
        if not fix: self.species_variable.append(s)

    def addOneReaction(self, reaction):
        r = Reaction(reaction)
        r.addFormula(self.ndic[r.name]['RateEq'].replace('self.',''))
        if self.ndic[r.name]['Type'] == 'Irrev': r.reversible = False

        #fxnames = self.hasFixedSpecies()
        processed_parameter = []
        # where parameters are defined `locally' per reaction
        for p in self.ndic[r.name]['Params']:
            p = p.replace('self.','')
            if p not in self.hasGlobalParameters() and p not in self.hasFixedSpecies():
                print "\tAdding parameter %s from networkdict" % p
                self.addParameter(p)
                par = self.__getattribute__(p)
                par.setAssociation(r)
                r.addParameter(par)
                processed_parameter.append(p)
            elif p not in self.hasFixedSpecies():
                print "\tUpdating parameter %s from networkdict" % p
                pidx = self.hasGlobalParameters().index(p)
                self.global_parameters[pidx].setAssociation(r)
                r.addParameter(self.global_parameters[pidx])
                processed_parameter.append(p)
        #print self.hasGlobalParameters()
        # where parameters are not `locally' defined and are extracted from Req (ie from SBML)
        for p in r._names:
            p = p.replace('self.','')
            if p not in self.hasGlobalParameters() and p not in self.hasSpecies():
                print "\tAdding parameter %s from formula" % p
                self.addParameter(p)
                par = self.__getattribute__(p)
                par.setAssociation(r)
                r.addParameter(par)
                processed_parameter.append(p)
            elif p not in processed_parameter and p not in self.hasSpecies():
                print "\tUpdating parameter %s from formula" % p
                pidx = self.hasGlobalParameters().index(p)
                self.global_parameters[pidx].setAssociation(r)
                r.addParameter(self.global_parameters[pidx])
                processed_parameter.append(p)
                
        self.__setattr__(reaction, r)
        self.reactions.append(r)

    def addParameter(self, name):
        if self.__InitDict__.has_key(name):
            par = Parameter(name, self.__InitDict__[name])
        else:
            par = Parameter(name, 0.0)
            if name not in self.__not_inited__: self.__not_inited__.append(name)
        self.global_parameters.append(par)
        self.__setattr__(name, par)

    def addSpecies(self):
        self.species = []
        self.species_variable = []
        for s in self.__model__.species:
            if self.__InitDict__.has_key(s):
                self.addOneSpecies(s, self.__InitDict__[s], fix=False)
            else:
                self.addOneSpecies(s, 0.0, fix=False)
        for s in self.__model__.fixed_species:
            if self.__InitDict__.has_key(s):
                self.addOneSpecies(s, self.__InitDict__[s], fix=True)
            else:
                self.addOneSpecies(s, 0.0, fix=True)


    def addReactions(self):
        self.reactions = []
        for r in self.__model__.reactions:
            self.addOneReaction(r)
        non_parameters = self.hasGlobalParameters()+self.hasSpecies()+self.hasFixedSpecies()
        for k in self.__InitDict__.keys():
            if k not in non_parameters:
                print 'Adding new parameter:', k
                self.addParameter(k)

    def replaceParameterWithRule(self, ar):
        par = self.__getattribute__(ar.name)
        for r in par.association:
            ar.setAssociation(r)
            setattr(r, ar.name, ar)
            r.parameters[r.hasParameters().index(ar.name)] = ar
        self.global_parameters[self.hasGlobalParameters().index(ar.name)] = ar
        self.__setattr__(ar.name, ar)

    def setAssignmentRules(self):
        for p in self.global_parameters:
            #print p.name
            if self.__forcing_function_dic__.has_key(p.name):
                print 'Assigning: %s = %s' % (p.name, self.__forcing_function_dic__[p.name])
                p2 = AssignmentRule(p.name, None)
                self.replaceParameterWithRule(p2)
                assert isinstance(p2, AssignmentRule), "\nHappy assertion error"
                #print type(p2)
                p2.addFormula(self.__forcing_function_dic__[p.name])
                print p2._names
                for n in p2._names:
                    p2.addModelAttr(self.__getattribute__(n))
                self.forcing_functions.append(p2)
                if p2.name in self.__not_inited__:
                    self.__not_inited__.pop(self.__not_inited__.index(p.name))

    def generateMappings(self):
        for reac in self.reactions:
            for reag in self.ndic[reac.name]['Reagents']:
                if self.ndic[reac.name]['Reagents'][reag] < 0.0:
                    reac.addSubstrate(self.__getattribute__(reag.replace('self.','')))
                    self.__getattribute__(reag.replace('self.','')).setSubstrate(self.__getattribute__(reac.name))
                else:
                    reac.addProduct(self.__getattribute__(reag.replace('self.','')))
                    self.__getattribute__(reag.replace('self.','')).setProduct(self.__getattribute__(reac.name))
                reac.stoichiometry.setdefault(reag.replace('self.',''), self.ndic[reac.name]['Reagents'][reag])
            for mod in self.ndic[reac.name]['Modifiers']:
                reac.addModifier(self.__getattribute__(mod.replace('self.','')))
                self.__getattribute__(mod.replace('self.','')).setModifier(self.__getattribute__(reac.name))

    def setStoichiometricMatrix(self):
        vspec = self.hasVariableSpecies()
        react = self.hasReactions()
        nm = numpy.zeros((len(vspec), len(react)),'d')
        for sp in vspec:
            for r in self.get(sp).isReagentOf():
                nm[vspec.index(sp)][react.index(r)] = self.get(r).stoichiometry[sp]
            # this is if absolute stoichiometry value is used
            ##  for r in self.get(sp).isSubstrateOf():
                ##  nm[vspec.index(sp)][react.index(r)] = abs(self.get(r).stoichiometry[sp])
            ##  for r in self.get(sp).isProductOf():
                ##  nm[vspec.index(sp)][react.index(r)] = -abs(self.get(r).stoichiometry[sp])
        self.stoichiometric_matrix = StructMatrix(nm, range(len(vspec)), range(len(react)))
        self.stoichiometric_matrix.setRow(vspec)
        self.stoichiometric_matrix.setCol(react)

    def addODEs(self):
        self.ODEs = []
        for varspec in self.stoichiometric_matrix.row:
            if self.struct != None:
                if varspec not in self.struct.Nr.row:
                    print 'Creating dependent ODE_%s' % varspec
                    ode = ODE(self.get(varspec), independent=False)
                else:
                    print 'Creating independent ODE_%s' % varspec
                    ode = ODE(self.get(varspec), independent=True)
            else:
                print 'Creating independent* ODE_%s (*assumed - no structural information available)' % varspec
                ode = ODE(self.get(varspec), independent=True)
            mrow = self.stoichiometric_matrix.getRowsByName(varspec)
            for e in range(len(mrow[0])):
                if mrow[0,e] != 0.0:
                    print '\tAdding term: %s*%s' % (mrow[0,e], self.stoichiometric_matrix.col[e])
                    ode.addReaction(self.get(self.stoichiometric_matrix.col[e]), mrow[0,e])
            self.__setattr__(ode.name, ode)
            self.ODEs.append(ode)
            self.__setattr__('xcode_'+ode.name, compile(ode.getGlobalFormula(), '<string>', 'exec'))
            
    def hasODEs(self):
        return MapList([o.name for o in self.ODEs])

    def evalODEs(self, odes):
        return [v() for v in odes]
        
    def evalXcode(self, ode):
        exec(self.__getattribute__('xcode_'+ode.name))
        return sdot

    def hasReactions(self):
        return MapList([r.name for r in self.reactions])

    def hasSpecies(self):
        return MapList([s.name for s in self.species])

    def hasFixedSpecies(self):
        return MapList([s.name for s in self.species if s.fixed])

    def hasVariableSpecies(self):
        return MapList([s.name for s in self.species if not s.fixed])

    def findReactionsThatIncludeAllSpecifiedReagents(self, *args):
        assert len(args) > 1, '\nNeed two or more species for this one!'
        setlist = [self.__getattribute__(s).isReagentOf().asSet() for s in args]
        isect = setlist[0]
        for s in setlist:
            isect.intersection_update(s)
        return MapList(isect)

    def hasGlobalParameters(self):
        return MapList(p.name for p in self.global_parameters)

##  if __psyco_active__:
    ##  psyco.bind(NewCoreBase)
    ##  psyco.bind(NumberBase)
    ##  psyco.bind(Species)
    ##  psyco.bind(Parameter)
    ##  psyco.bind(AssignmentRule)
    ##  psyco.bind(Reaction)
    ##  psyco.bind(ODE)
    ##  psyco.bind(NewCore)

            

if __name__ == '__main__':
    
    print "\n*****\nTESTING NEWCORE\n*****\n"


    import pysces
    M = pysces.model('isola_thesis')
    M.doLoad()

    print '\nModel', M.ModelFile
    print '============='
    modmap = NewCore(M)

    print 'Reactions\n', modmap.hasReactions()
    print 'Species\n', modmap.hasSpecies()
    print 'FixedSpecies\n', modmap.hasFixedSpecies()
    print ' '
    print 'R1 has reagents\n', modmap.R1.hasReagents()
    print 'R1 has sub\n', modmap.R1.hasSubstrates()
    print 'R1 has prod\n', modmap.R1.hasProducts()
    print 'R1 has mod\n', modmap.R1.hasModifiers()
    print ' '
    print 'P is reagent\n', modmap.P.isReagentOf()
    print 'P is sub\n', modmap.P.isSubstrateOf()
    print 'P is prod\n', modmap.P.isProductOf()
    print 'P is mod\n', modmap.P.isModifierOf()
    print ' '
    print 'R2 stoich\n', modmap.R2.stoichiometry
    print ' '
    print 'findReactionsThatIncludeAllSpecifiedReagents(A, B):', modmap.findReactionsThatIncludeAllSpecifiedReagents('A','B')


    print '\nmodmap.hasGlobalParameters\n', modmap.hasGlobalParameters()

    print '\nParameter associations'
    for p in modmap.global_parameters:
        print '%s.isParameterOf() %s' % (p.name, p.isParameterOf())

    print "\n*****\nTESTING INFIXPARSER\n*****\n"

    # Build the lexer and try it out
    ##  m = MyInfixLexer()
    m = MyInfixParser()
    m.buildlexer()           # Build the lexer
    ##  m.testlexer("(K3*K4+s0)**(h)")     # Test it
    m.buildparser(debug=0, debugfile='infix.dbg', tabmodule='infix_tabmodule')
    m.setNameStr('self.', '()')
    print "\ngo!\n"
    
    m.parse("(K3*K4+s0)**(h)")
    print '\nm.input\n', m.input
    print 'm.output\n', m.output
    print 'm.names\n', m.names
    
    m.parse("(2**K3*K4+s0)**(h/5)")
    print '\nm.input\n', m.input
    print 'm.output\n', m.output
    print 'm.names\n', m.names